#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@Project ：python_learning 
@File ：共享内存操作.py
@IDE  ：PyCharm 
@Author ：李涵彬
@Date ：2024/12/29 12:56 
"""

from multiprocessing import shared_memory
from typing import Any, Dict


class SharedMemoryManager:
	def __init__(self):
		"""
		初始化共享内存管理器。
		"""
		self.shared_memories: Dict[str, shared_memory.SharedMemory] = {}

	def create_shared_memory(self, name: str, size: int):
		"""
		创建一块共享内存。

		:param name: 共享内存的名称。
		:param size: 共享内存的大小（字节）。
		"""
		try:
			self.shared_memories[name] = shared_memory.SharedMemory(name=name, create=True, size=size)
			print(f"创建了名为 {name} 的共享内存，大小为 {size} 字节。")
		except Exception as e:
			print(f"创建共享内存失败：{e}")

	def open_existing_shared_memory(self, name: str):
		"""
		打开一个已经存在的共享内存。

		:param name: 共享内存的名称。
		"""
		try:
			self.shared_memories[name] = shared_memory.SharedMemory(name=name)
			print(f"打开了名为 {name} 的共享内存。")
		except Exception as e:
			print(f"打开共享内存失败：{e}")

	def set_value(self, name: str, value: Any, offset: int = 0):
		"""
		设置共享内存中的值。

		:param name: 共享内存的名称。
		:param value: 要设置的值。
		:param offset: 偏移量，默认为0。
		"""
		try:
			shared_mem = self.shared_memories[name]
			# 将值转换为字节，然后复制到共享内存的指定位置
			shared_mem.buf[offset:offset + len(value)] = value.encode()
			print(f"在 {name} 的共享内存中，偏移量 {offset} 处设置值为 {value}。")
		except Exception as e:
			print(f"设置共享内存值失败：{e}")

	def get_value(self, name: str, size: int = 10, offset: int = 0) -> Any:
		"""
		获取共享内存中的值。

		:param name: 共享内存的名称。
		:param size: 要获取的值的大小，默认为10字节。
		:param offset: 偏移量，默认为0。
		"""
		try:
			shared_mem = self.shared_memories[name]
			# 从共享内存的指定位置读取指定大小的数据
			value = shared_mem.buf[offset:offset + size].tobytes().decode()
			print(f"从 {name} 的共享内存中，偏移量 {offset} 处获取值为 {value}。")
			return value
		except Exception as e:
			print(f"获取共享内存值失败：{e}")

	def close_shared_memory(self, name: str):
		"""
		关闭共享内存。

		:param name: 共享内存的名称。
		"""
		try:
			shared_mem = self.shared_memories[name]
			shared_mem.close()
			print(f"关闭了名为 {name} 的共享内存。")
		except Exception as e:
			print(f"关闭共享内存失败：{e}")

	def unlink_shared_memory(self, name: str):
		"""
		释放共享内存。

		:param name: 共享内存的名称。
		"""
		try:
			shared_mem = self.shared_memories[name]
			shared_mem.unlink()
			print(f"释放了名为 {name} 的共享内存。")
		except Exception as e:
			print(f"释放共享内存失败：{e}")

	def __enter__(self):
		"""
		支持上下文管理器。
		"""
		return self

	def __exit__(self, exc_type, exc_val, exc_tb):
		"""
		上下文管理器退出时清理资源。
		"""
		for name in list(self.shared_memories.keys()):
			self.unlink_shared_memory(name)


# 示例用法：
if __name__ == "__main__":
	with SharedMemoryManager() as shm:
		shm.create_shared_memory("example", 1024)
		shm.set_value("example", "Hello, Shared Memory!")
		value = shm.get_value("example", len("Hello, Shared Memory!"))  # 确保大小足以容纳字符串
		print(f"从共享内存中获取的值为：{value}")
		shm.close_shared_memory("example")
		shm.unlink_shared_memory("example")
